# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *

import datetime
import numpy as np
import pandas as pd
import multiprocessing
import statsmodels.api as sm


def init(context):
    # 目标指数
    context.index_code = 'SHSE.000852'# 中证1000 'SHSE.000852'， 沪深300 'SHSE.000300'
    # 因子计算长度
    context.urt_length = 20# 20天交易日
    # 持仓数量
    context.holding_num = 30
    # 定时任务，日频，月底换仓的逻辑在algo中进行判断
    schedule(schedule_func=algo, date_rule='1d', time_rule='14:50:00')


def algo(context):
    # 当前时间
    now_date = context.now.strftime('%Y-%m-%d')

    next_day = get_N_trading_date(context,now_date,counts=1,model='next',exchange='SHSE')
    # 判断是否为每个月最后一个交易日
    if context.now.month!=pd.Timestamp(next_day).month:
        # 获取上一个交易日日期
        last_date = get_N_trading_date(context,now_date,counts=1,model='previous',exchange='SHSE')

        # 获取指数成分股
        stocks = list(stk_get_index_constituents(index=context.index_code, trade_date=last_date)['symbol'])

        # 计算传统换手率并去极值、标准化、市值中性化
        turnrate = stk_get_daily_basic_pt_new(context, symbols=stocks, fields='turnrate', end_date=last_date, counts=context.urt_length, df=True)
        turnrate20 = turnrate.groupby('symbol')['turnrate'].mean()
        turnrate20 = winsorize_med(turnrate20)
        turnrate20 = standardlize(turnrate20)
        turnrate20 = neutralize_MarketValue(context, data=turnrate20, date=last_date, counts=1)

        # 计算量稳换手率并去极值、标准化、市值中性化
        turnrate_STR = turnrate.groupby('symbol')['turnrate'].std()
        turnrate_STR = winsorize_med(turnrate_STR)
        turnrate_STR = standardlize(turnrate_STR)
        turnrate_STR = neutralize_MarketValue(context, data=turnrate_STR, date=last_date, counts=1)
        
        # 计算UTR2.0换手率
        turnrate_UTR2 = cal_factor_UTR2(turnrate20,turnrate_STR).sort_values()
        to_buy = list(turnrate_UTR2[:context.holding_num].index)
        print('{}:待买入股票池{}只：{}'.format(context.now,len(to_buy),to_buy))

        # 交易
        positions = context.account().positions()
        holding_symbols = [posi['symbol'] for posi in positions]
        new_prices = history(symbol=holding_symbols, frequency='1d', start_time=now_date, end_time=now_date, fields='symbol,close', adjust=ADJUST_PREV, adjust_end_time=context.backtest_end_time, df=False)
        # 平不在标的池的股票（注：日频市价单，backtest_match_mode=1，以当日收盘价进行交易）
        for position in positions:
            symbol = position['symbol']
            if symbol not in to_buy:
                new_price = list(filter(lambda x:x[r'symbol'] == symbol, new_prices))
                if len(new_price)==0:continue
                order_target_percent(symbol=symbol, percent=0, order_type=OrderType_Market,position_side=PositionSide_Long,price=new_price[0]['close']*0.99)

        # 获取股票的权重
        percent = 1 / len(to_buy)
        # 买在标的池中的股票（注：日频市价单，backtest_match_mode=1，以当日收盘价进行交易）   
        new_prices = history(symbol=to_buy, frequency='1d', start_time=now_date, end_time=now_date, fields='symbol,close', adjust=ADJUST_PREV, adjust_end_time=context.backtest_end_time, df=False)
        for symbol in to_buy:
            # 收盘价（日频数据）
            new_price = list(filter(lambda x:x[r'symbol'] == symbol, new_prices))
            if len(new_price)==0:continue
            order_target_percent(symbol=symbol, percent=percent, order_type=OrderType_Market,position_side=PositionSide_Long,price=new_price[0]['close']*1.01)


def cal_factor_UTR2(turn20,STR):
    """
    UTR2.0 = STR + 𝑠𝑜𝑓𝑡𝑠𝑖𝑔𝑛(STR) ∙ Turn20
    """
    return STR+np.multiply(softsign(STR),turn20)


def 𝑠𝑜𝑓𝑡𝑠𝑖𝑔𝑛(x):
    return x/(1+abs(x))


def stk_get_daily_basic_pt_new(context, symbols, fields, start_date=None, end_date=None, counts:int = None, df=False):
    """
    多日期调用stk_get_daily_basic_pt函数，当有count时就不采用start_date,count为正整数
    """
    if counts!=None:
        date_list = get_N_trading_date(context,date=end_date,counts=counts,model='previous',exchange='SHSE',return_list=True)
    else:
        date_list = get_trading_dates(exchange='SZSE', start_date=start_date, end_date=end_date)
    
    # 循环获取数据
    df_total = pd.DataFrame()
    for date in date_list:
        df_new = stk_get_daily_basic_pt(symbols=symbols, fields=fields, trade_date=date, df=df)
        if df:
            df_total = pd.concat([df_total,df_new])
        else:
            df_total = df_total+df_new
    return df_total


def stk_get_daily_mktvalue_pt_new(context, symbols, fields, start_date=None, end_date=None, counts:int = None, df=False):
    """
    多日期调用stk_get_daily_mktvalue_pt函数，当有count时就不采用start_date,count为正整数
    """
    if counts!=None:
        date_list = get_N_trading_date(context,date=end_date,counts=counts,model='previous',exchange='SHSE',return_list=True)
    else:
        date_list = get_trading_dates(exchange='SZSE', start_date=start_date, end_date=end_date)
    
    # 循环获取数据
    df_total = pd.DataFrame()
    for date in date_list:
        df_new = stk_get_daily_mktvalue_pt(symbols=symbols, fields=fields, trade_date=date, df=df)
        if df:
            df_total = pd.concat([df_total,df_new])
        else:
            df_total = df_total+df_new
    return df_total


def winsorize_med(data, scale=3, inclusive=True, inf2nan=True):
    """
    去极值
    :param data：待处理数据[Series]
    :param scale：标准差倍数，默认为3
    :param inclusive：True为将边界外的数值调整为边界值，False为将边界外的数值调整为NaN
    :param inf2nan：True为将inf转化为nan，False不转化
    """
    data = data.astype('float')
    if inf2nan:
        data = data.replace([np.inf, -np.inf], np.nan)
    std_ = data.std()
    mean_ = data.mean()
    if inclusive:
        data[data>mean_+std_*scale]=mean_+std_*scale
        data[data<mean_-std_*scale]=mean_-std_*scale
    else:
        data[data>mean_+std_*scale]=np.nan
        data[data<mean_-std_*scale]=np.nan
    return data


def standardlize(data, inf2nan=True):
    """
    标准化
    :param data:待处理数据
    :param inf2nan：是否将inf转化为nan
    """
    if inf2nan:
        data = data.replace([np.inf, -np.inf], np.nan)
    return (data - data.mean()) / data.std()


def neutralize_MarketValue(context,data,date,counts=1):
    """
    市值中性化
    :param data:待处理数据
    :param date:目标日期
    :param counts：历史回溯天数
    """
    if isinstance(data,pd.Series):
        data = data.to_frame()
    security = data.index.to_list()
    market_value = stk_get_daily_mktvalue_pt_new(context, symbols=security, fields='tot_mv', end_date=date, counts=counts, df=True)
    market_value_mean = market_value.groupby('symbol')['tot_mv'].mean()
    x = sm.add_constant(market_value_mean)
    common_index = list(set(x.index) & set(data.index))
    x = x.loc[common_index,:]
    data = data.loc[common_index,:]
    residual = sm.OLS(data, x).fit().resid# 此处使用最小二乘回归计算
    return residual


def get_N_trading_date(context,date,counts=1,model='previous',exchange='SHSE',return_list=False):
    """
    获取end_date前N个交易日,end_date为datetime格式，不包括date日期
    :param date：目标日期
    :param counts：历史回溯天数，默认为1，即前一天，该值需为非零正整数
    :param model: 模式，默认为previous，即前N个交易日；next模式为后N个交易日
    :return_list: 是否返回日期列表，默认为False
    """
    if counts<=0 or not isinstance(counts, int):
        print('参数类型错误：get_N_trading_date()的counts参数需为非零正整数.')
        import sys
        sys.exit(-1)
    date = pd.Timestamp(date)
    date_str = date.strftime('%Y-%m-%d')
    # 判断context.trading_dates是否存在
    if 'trading_dates' not in {k: v for k, v in context.__dict__.items() if not k.startswith('__')}.keys():
        context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=date.year-1, end_year=datetime.datetime.now().year)
    # 判断当前日期是否为交易日
    pre_times = 1
    if date_str not in list(context.trading_dates['trade_date']) and str(date.year) in set([x[:4] for x in context.trading_dates['trade_date']]):# 非交易日
        date = context.trading_dates[context.trading_dates['date']==date_str]['pre_trade_date'].iloc[0]
        pre_times = 0
    date = pd.Timestamp(date)
    date_str = date.strftime('%Y-%m-%d')
    # 历史N个交易日
    for i in range(2):
        trading_dates_dropna = context.trading_dates.replace('',np.nan).dropna(how='any')
        if model == 'previous':
            previous_date_df = trading_dates_dropna[trading_dates_dropna['trade_date']<=date_str]
            if len(previous_date_df)>=counts+1:
                if return_list:
                    return list(previous_date_df['trade_date'].iloc[-counts-pre_times:])
                else:
                    return previous_date_df['trade_date'].iloc[-counts-pre_times]
            else:
                start_date = pd.Timestamp(context.trading_dates['date'].iloc[0])-datetime.timedelta(days = int(365*np.ceil((counts-len(previous_date_df))/250)))
                start_year = min(start_date.year-1,date.year)
                end_year = int(context.trading_dates['date'].iloc[-1][:4])
                context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=start_year, end_year=end_year)
        elif model == 'next':
            previous_date_df = trading_dates_dropna[trading_dates_dropna['trade_date']>=date_str]
            if len(previous_date_df)>=counts+1:
                if return_list:
                    return list(previous_date_df['trade_date'].iloc[:counts])
                else:
                    return previous_date_df['trade_date'].iloc[counts]
            else:
                end_date = pd.Timestamp(context.trading_dates['date'].iloc[-1])+datetime.timedelta(days = int(365*np.ceil((counts-len(previous_date_df))/250)))
                start_year = int(context.trading_dates['date'].iloc[0][:4])
                end_year = max(end_date.year,date.year)
                context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=start_year, end_year=end_year)


# 获取每次回测的报告数据
def on_backtest_finished(context, indicator):
    print('回测结束。')
    print(indicator)


if __name__ == '__main__':
    '''
        strategy_id策略ID,由系统生成
        filename文件名,请与本文件名保持一致
        mode实时模式:MODE_LIVE回测模式:MODE_BACKTEST
        token绑定计算机的ID,可在系统设置-密钥管理中生成
        backtest_start_time回测开始时间
        backtest_end_time回测结束时间
        backtest_adjust股票复权方式不复权:ADJUST_NONE前复权:ADJUST_PREV后复权:ADJUST_POST
        backtest_initial_cash回测初始资金
        backtest_commission_ratio回测佣金比例
        backtest_slippage_ratio回测滑点比例
        backtest_match_mode市价撮合模式，以下一tick/bar开盘价撮合:0，以当前tick/bar收盘价撮合：1
    '''
    run(strategy_id='859c3b0a-1765-11ec-acc4-d0509949eed8',
        filename='main.py',
        mode=MODE_BACKTEST,
        token='47ca47f849b3a0f66ec0f7013bb56bb667d63a70',
        backtest_start_time='2018-12-28 08:00:00',
        backtest_end_time='2023-10-08 16:00:00',
        backtest_adjust=ADJUST_PREV,
        backtest_initial_cash=1000000,
        backtest_commission_ratio=0.0007,
        backtest_slippage_ratio=0.00123,
        backtest_match_mode=1)